#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time         : 19-2-9 下午12:48
# @Author       : ding
# @File         : tfread.py
# @Description  : tfrecords 读取基类

import tensorflow as tf
from abc import ABC, abstractmethod

from PIL import Image


class TFReader(ABC):
    def __init__(self, train_shape, nthread=2, shuffle=False, numb_epochs=1):
        self.train_shape = train_shape
        self.nthread = nthread
        self.shuffle = shuffle
        self.numb_epochs = numb_epochs

    @staticmethod
    def init(sess):
        g_init = tf.global_variables_initializer()
        l_init = tf.local_variables_initializer()
        sess.run([g_init, l_init])

    @abstractmethod
    def improc(self, example, args=()):
        # image = example[0]
        # image = tf.image.resize_images(image, self.train_shape[:2])
        # image = tf.reshape(image, self.train_shape)
        # example[0] = image
        return example

    @abstractmethod
    def feature_handles(self, serialized_example):
        # features = tf.parse_single_example(serialized_example, features={
        #     "img_raw": tf.FixedLenFeature([], tf.string),
        #     "img_lab": tf.FixedLenFeature([], tf.int64),
        #     "img_shape": tf.FixedLenFeature([3], tf.int64)})
        #
        # img_raw = tf.decode_raw(features["img_raw"], tf.uint8)
        # img_lab = tf.cast(features["img_lab"], tf.int64)
        # img_shape = tf.cast(features["img_shape"], tf.int32)
        # image = tf.reshape(img_raw, img_shape)
        # label = tf.reshape(img_lab, [1])

        return []

    def read_turn(self, file_names, batch_size, args=()):
        if not isinstance(file_names, list):
            file_names = [file_names]
        filename_queue = tf.train.string_input_producer(
            file_names, num_epochs=self.numb_epochs, shuffle=self.shuffle)
        reader = tf.TFRecordReader()
        _, serialized_example = reader.read(filename_queue)

        example = self.feature_handles(serialized_example)

        news = self.improc(example, *args)

        batch_data = tf.train.batch(list(news),
                                    batch_size=batch_size,
                                    num_threads=self.nthread,
                                    capacity=512+batch_size*2*self.nthread)
        return batch_data

    def read_shuf(self, file_names, batch_size, args=()):
        if not isinstance(file_names, list):
            file_names = [file_names]
        filename_queue = tf.train.string_input_producer(file_names, num_epochs=self.numb_epochs)
        reader = tf.TFRecordReader()
        _, serialized_example = reader.read(filename_queue)

        example = self.feature_handles(serialized_example)

        news = self.improc(example, *args)

        batch_data = tf.train.shuffle_batch(list(news),
                                            batch_size=batch_size,
                                            num_threads=self.nthread,
                                            capacity=512+batch_size*2*self.nthread,
                                            min_after_dequeue=batch_size)
        return batch_data


class TFNormalReader(TFReader):
    def __init__(self, train_shape, nthread=2, shuffe=False, numb_epochs=1, one_hot=0):
        TFReader.__init__(self, train_shape, nthread, shuffe, numb_epochs)
        self.one_hot = one_hot

    def feature_handles(self, serialized_example):
        features = tf.parse_single_example(serialized_example, features={
            "image/encode": tf.FixedLenFeature([], tf.string),
            "image/label": tf.FixedLenFeature([], tf.int64),
            "image/shape": tf.FixedLenFeature([3], tf.int64)})

        img_raw = tf.decode_raw(features["image/encode"], tf.uint8)
        img_lab = tf.cast(features["image/label"], tf.int64)
        img_shape = tf.cast(features["image/shape"], tf.int32)
        image = tf.reshape(img_raw, img_shape)
        label = img_lab

        return [image, label]

    def improc(self, example, args=()):
        image, label = example
        image = tf.image.resize_images(image, self.train_shape[:2])
        image = tf.reshape(image, self.train_shape)
        image /= 255

        if self.one_hot > 1:
            label = tf.one_hot(label, self.one_hot)

        return [image, label]


if __name__ == "__main__":
    tfd = TFNormalReader([224, 448, 3], one_hot=10)
    train_batch, train_label_batch = tfd.read_shuf(["idcard_train_001_of_000.tfrecords",
                                                    "idcard_train_001_of_001.tfrecords"], 10)

    with tf.Session() as sess:
        tfd.init(sess)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        train_batch, train_label_batch = sess.run([train_batch, train_label_batch])
        print(train_batch.shape, type(train_batch[0]))

        image = Image.fromarray((255 * train_batch[3]).astype('uint8'))
        image.save("2.jpg")
        print(train_label_batch[0])

        coord.request_stop()
        coord.join(threads)
